import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, optimizers, activations, losses, metrics, \
    callbacks, utils
import sys
import os
from python_ai.common.xcommon import *
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt

np.random.seed(777)
tf.random.set_seed(777)
filename = os.path.basename(__file__)

LEN_DICT = 1000
N_STEPS = 80
N_EMBEDDING = 300
N_RNN_HIDDEN = 128

(x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data(num_words=LEN_DICT)
check_shape(x_train, 'x_train')  # (25000, )  # an ndarray from ragged nested sequences
check_shape(y_train, 'y_train')  # (25000, )
check_shape(x_test, 'x_test')  # (25000, )  # an ndarray from ragged nested sequences
check_shape(y_test, 'y_test')  # (25000, )

x_train = keras.preprocessing.sequence.pad_sequences(x_train, maxlen=N_STEPS)
x_test = keras.preprocessing.sequence.pad_sequences(x_test, maxlen=N_STEPS)
check_shape(x_train, 'x_train')  # (25000, 80)
check_shape(y_train, 'y_train')  # (25000, )
check_shape(x_test, 'x_test')  # (25000, 80)
check_shape(y_test, 'y_test')  # (25000, )

layer_embedding = layers.Embedding(LEN_DICT, N_EMBEDDING, input_length=N_STEPS)
x_embedded = layer_embedding(x_train)
check_shape(x_embedded, 'x_embedded')
x_out = x_embedded

layer_lstm01 = layers.SimpleRNN(N_RNN_HIDDEN, return_sequences=True, unroll=True, dropout=0.5)
x_out = layer_lstm01(x_out, training=True)
check_shape(x_out, 'x_out')

layer_lstm02 = layers.SimpleRNN(N_RNN_HIDDEN, return_sequences=False, unroll=True, dropout=0.5)
x_out = layer_lstm02(x_out, training=True)
check_shape(x_out, 'x_out')

layer_fc = layers.Dense(1, activation=activations.sigmoid)
x_out = layer_fc(x_out, training=True)
check_shape(x_out, 'x_out')

print('Studying over.')
